# Data normalisation methods for ConsensusCluster

"""

Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import numpy as N
import sys

def swrap(func):
    """Allow matrix normalisation functions to take sdata objects as input"""

    def wrapped(obj, **kwds):
        if isinstance(obj, N.ndarray):
            return func(obj, **kwds)
        
        else:
            try:
                obj.M = func(obj.M, **kwds)
            except AttributeError:
                raise ValueError, 'Matrix object must be ndarray or BaseParser'

    return wrapped
        
@swrap
def normalise(M, log2=True, sub_medians=True, center=True, scale=True):
    """
    Perform a number of normalisation routines on M
    
    log2                - log2 transform the data (Yes if this is raw gene data)
    sub_medians         - subtract the median of medians from each data value
    center              - subtract the data by its average, making the overall mean 0
    scale               - subtract the root-mean-square of the data AFTER centering

    """

    if log2:
        M = log2_transform(M)

    if sub_medians:
        M = subtract_medians(M)

    if scale or center:
        M = center_matrix(M)

    if scale:
        M = scale_matrix(M)

    return M

@swrap
def center_matrix(M):
    """Subtract the mean from matrix M, resulting in a matrix with mean 0"""

    return (M - N.average(M, 0))

@swrap
def scale_matrix(M):
    """Subtract the root-mean-square from each data member"""

    N.seterr(all='raise')   #Catch divide-by-zero, otherwise SVD won't converge

    T = N.transpose(M)
    for i in range(M.shape[1]):
        if T[i].any():
            T[i] = T[i] / N.sqrt(N.sum(T[i]**2) / (M.shape[0] - 1))

    return M

@swrap
def log2_transform(M):
    """Take the log2 of each value in M"""

    if not M.all():
        print('WARNING: Zero values in log transform! Using log2(M + 1) instead')
        return N.log2(M + 1)
    
    return N.log2(M)

@swrap
def subtract_medians(M):
    """Subtract each value in M by the median of medians"""

    return M - N.median(M)
    
@swrap
def subtract_feature_medians(M):
    """Subtract each column in M by the median over all rows"""

    return M - N.median(M, 0)

@swrap
def row_normalise_mean(M):
    """Normalise rows to mean 0"""

    return (M.T - N.average(M, 1)).T

@swrap
def row_normalise_median(M):
    """Normalise rows to median 0"""

    return (M.T - N.median(M, 1)).T
